//
//  BTLSSLSocketHandler.m
//  BTLSSLSocketHandlers
//
//  Created by Buckley on 12/28/08.
//  Copyright 2008 __MyCompanyName__. All rights reserved.
//

#import "BTLSSLSocketHandler.h"
#import <openssl/err.h>

#ifdef WIN32
#define pipe(fds) _pipe(fds, SSL_BUF_SIZE, _O_BINARY)
#endif

#define SSL_BUF_SIZE 65536
#define SSL_ERROR_BUF_SIZE 128

@interface BTLSSLSocketHandler (Private)
- (void)initPipes;
@end

@implementation BTLSSLSocketHandler

+ (void)load
{
	SSL_library_init();
	SSL_load_error_strings();
	ERR_load_BIO_strings();
	OpenSSL_add_all_algorithms();
}

- (id)initWithMethod:(SSL_METHOD*)aMethod
{
	id ret = [super init];
	if(ret != nil){
		self = ret;
		method = aMethod;
		context = SSL_CTX_new(method);
		ssl = SSL_new(context);
	}
	
	return ret;
}

- (void)dealloc
{
	if([self encapsulatingList] != nil){
		SSL_shutdown(ssl);
		int readBytes = 0;
		char writeBuffer[SSL_BUF_SIZE];
		BTLSocketBuffer* negotiationBuffer = [BTLSocketBuffer new];
		do{
			readBytes = read(remoteReadfd, writeBuffer, SSL_BUF_SIZE);
			[negotiationBuffer addData:writeBuffer ofSize:readBytes];
		}while(readBytes > 0);
		[self finishedWritingData:negotiationBuffer toAddress:[[[self encapsulatingList] socket] remoteAddress]];
	}
	
	close(readfd);
	close(writefd);
	close(remoteReadfd);
	close(remoteWritefd);
	[super dealloc];
}

- (id)copyWithZone:(NSZone*)zone
{
	BTLSSLSocketHandler* ret = [[[self class] alloc] initWithMethod:method];
	if(ret != nil){
		[ret initPipes];
		ret->method = method;
		ret->context = context;
		ret->ssl = ssl;
	}
	return ret;
}

- (void)readData:(BTLSocketBuffer*)someData fromAddress:(BTLSocketAddress*)anAddress
{
	const char* cipherText = [someData rawData];
	uint64_t remainingCipherText = [someData size];
	BTLSocketBuffer* plainText = [BTLSocketBuffer new];
	char readBuffer[SSL_BUF_SIZE];
	while(remainingCipherText > 0){
		int writtenBytes = write(remoteWritefd, cipherText, remainingCipherText);
		if(writtenBytes > 0){
			remainingCipherText -= writtenBytes;
		}else if(writtenBytes != EAGAIN){
			break;
		}
		
		int readBytes = 0;
		do{
			readBytes = SSL_read(ssl, readBuffer, SSL_BUF_SIZE);
			if(readBytes == 0){
				id delegate = [[[self encapsulatingList] socket] delegate];
				if([delegate respondsToSelector:@selector(sslConnectionClosedByRemoteHostOnHandler:)]){
					[delegate sslConnectionClosedByRemoteHostOnHandler:self];
				}
				[self removeSelf];
			}else if(readBytes < 0){
				if(SSL_get_error(ssl, readBytes) == SSL_ERROR_WANT_WRITE){
					char writeBuffer[SSL_BUF_SIZE];
					BTLSocketBuffer* negotiationBuffer = [BTLSocketBuffer new];
					do{
						readBytes = read(remoteReadfd, writeBuffer, SSL_BUF_SIZE);
						[negotiationBuffer addData:writeBuffer ofSize:readBytes];
					}while(readBytes > 0);
					[self finishedWritingData:negotiationBuffer toAddress:anAddress];
				}else{
					readBytes = 0;
				}
			}else{
				[plainText addData:readBuffer ofSize:readBytes];
			}
		}while(readBytes > 0);
	}
	
	[self finishedReadingData:plainText fromAddress:anAddress];
}

- (BOOL)writeData:(BTLSocketBuffer*)someData toAddress:(BTLSocketAddress*)anAddress
{
	const char* plainText = [someData rawData];
	uint64_t remainingPlainText = [someData size];
	BTLSocketBuffer* cipherText = [BTLSocketBuffer new];
	char writeBuffer[SSL_BUF_SIZE];
	while(remainingPlainText > 0){
		int writtenBytes = SSL_write(ssl, plainText, remainingPlainText);
		if(writtenBytes > 0){
			remainingPlainText -= writtenBytes;
		}else{
			switch(SSL_get_error(ssl, writtenBytes)){
				case SSL_ERROR_WANT_WRITE:
				{
					char writeBuffer[SSL_BUF_SIZE];
					BTLSocketBuffer* negotiationBuffer = [BTLSocketBuffer new];
					do{
						writtenBytes = read(remoteReadfd, writeBuffer, SSL_BUF_SIZE);
						[negotiationBuffer addData:writeBuffer ofSize:writtenBytes];
						[self finishedWritingData:negotiationBuffer toAddress:anAddress];
					}while(writtenBytes > 0);
				}
					break;
					
				case SSL_ERROR_WANT_READ:
				{
					[[[self encapsulatingList] socket] read];
				}
					break;
				default:
					break;
			}
		}
		
		int readBytes = 0;
		do{
			readBytes = read(remoteReadfd, writeBuffer, SSL_BUF_SIZE);
			[cipherText addData:writeBuffer ofSize:readBytes];
		}while(readBytes > 0);
	}
	
	[self finishedWritingData:cipherText toAddress:anAddress];
	return YES;
}

- (void)useX509Certificate:(X509*)aCertificate
{
	SSL_use_certificate(ssl, aCertificate);
}

- (void)useASN1Certificate:(unsigned char*)aCertificate ofLength:(int)aLength
{
	SSL_use_certificate_ASN1(ssl, aCertificate, aLength);
}

- (void)useCertificateFile:(const char*)aFile ofType:(int)aType
{
	SSL_use_certificate_file(ssl, aFile, aType);
}

- (void)usePrivateKey:(EVP_PKEY*)aKey
{
	SSL_use_PrivateKey(ssl, aKey);
}

- (void)useASN1PrivateKey:(unsigned char*)aKey ofType:(int)aType length:(int)aLength
{
	SSL_use_PrivateKey_ASN1(aType, ssl, aKey, aLength);
}

- (void)usePrivateKeyFile:(const char*)aFile ofType:(int)aType
{
	SSL_use_PrivateKey_file(ssl, aFile, aType);
}

- (void)useRSAPrivateKey:(RSA*)anRSA
{
	SSL_use_RSAPrivateKey(ssl, anRSA);
}

- (void)useASN1RSAPrivateKey:(unsigned char*)aKey ofLength:(int)aLength
{
	SSL_use_RSAPrivateKey_ASN1(ssl, aKey, aLength);
}

- (void)useRSAPrivateKeyFile:(const char*)aFile ofType:(int)aType
{
	SSL_use_RSAPrivateKey_file(ssl, aFile, aType);
}

- (int)checkPrivateKey
{
	return SSL_check_private_key(ssl);
}

@end

@implementation BTLSSLSocketHandler (Private)
- (void)initPipes
{
	int fds[2];
	pipe(&fds);
	readfd = fds[0];
	remoteWritefd = fds[1];
	
	int opts = fcntl(fds[0], F_GETFL, 0);
	if(opts >= 0 && (opts & O_NONBLOCK) != 0){
		opts = (opts ^ O_NONBLOCK);
		fcntl(fds[0], F_SETFL, opts);
	}
	
	opts = fcntl(fds[1], F_GETFL, 0);
	if(opts >= 0 && (opts & O_NONBLOCK) != 0){
		opts = (opts ^ O_NONBLOCK);
		fcntl(fds[1], F_SETFL, opts);
	}
	
	pipe(&fds);
	remoteReadfd = fds[0];
	writefd = fds[1];
	
	opts = fcntl(fds[0], F_GETFL, 0);
	if(opts >= 0 && (opts & O_NONBLOCK) != 0){
		opts = (opts ^ O_NONBLOCK);
		fcntl(fds[0], F_SETFL, opts);
	}
	
	opts = fcntl(fds[1], F_GETFL, 0);
	if(opts >= 0 && (opts & O_NONBLOCK) != 0){
		opts = (opts ^ O_NONBLOCK);
		fcntl(fds[1], F_SETFL, opts);
	}
	
	SSL_set_rfd(ssl, readfd);
	SSL_set_wfd(ssl, writefd);
}
@end
